"""
    Title

    author: wxz
    date: 
    github: https://github.com/xinzwang
"""

import numpy as np


def fun1(x, y):
    out = np.square(x) + 2 * np.square(y) - 4 * x * y - 0.5 * y
    return out


def fun2(x, y):
    out = np.square(x) + np.square(y)
    return out


def fun3(x, y):
    """
    min: -1.4118161227972283
    """
    out = np.sin(x) + np.sin(y) + 0.2 * (np.abs(x + y) + np.abs(x - y))
    return out


class F_Elli(object):
    """
    论文中, 待优化函数1
    """

    def __init__(self):
        # vec应为n维正交向量, 组成空间基底
        self.n = 2
        # self.vec = np.array([[1, 0, 0, 0],
        #                      [0, 1, 0, 0],
        #                      [0, 0, 1, 0],
        #                      [0, 0, 0, 1]])
        self.vec = np.array([[1, 0],
                             [0, 1]])

    def forward(self, x):
        """
        x: n维向量
        y: n维向量
        """
        n = self.n
        vec = self.vec
        out = np.zeros([n], dtype=np.float64)
        for i in range(n):
            out += np.power(1000, (i - 1) / (n - 1)) * x.dot(vec[i, :].T)
        return out


class F_Rosenbrock(object):
    """
    Rosenbrock函数    用于测试最优化算法性能的非凸函数
    f_min = f(1,1) = 0
    """

    def __init__(self, n=10):
        self.n = n

    def forward(self, x):
        """单个x计算"""
        a = x[0:self.n - 1]
        b = x[1:self.n]
        return np.sum(100 * (np.square(np.square(a) - b)) + np.square(a - 1))

    def forward_batch(self, x):
        """批量计算"""
        x = x.copy()  # 深拷贝
        if x.ndim == 1:
            return self.forward(x)
        N = x.shape[0]
        n = self.n
        out = np.zeros(N)
        for i in range(N):
            a = x[i, 0:n - 1]
            b = x[i, 1:n]
            out[i] = np.sum(100 * (np.square(np.square(a) - b)) + np.square(a - 1))
        return out


if __name__ == '__main__':
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib import cm
    from mpl_toolkits.mplot3d import Axes3D

    """
    绘制函数的三维图像
    """

    fig = plt.figure()

    x = np.arange(-20, 20, 1)
    y = np.arange(-1000, 1000, 1)
    x, y = np.meshgrid(x, y)

    # fun1 图像
    # z = fun1(x, y)

    # fun2 图像
    # z = fun2(x, y)

    # fun3 图像
    # z = fun3(x, y)

    f = F_Rosenbrock(2)
    z = f.forward([x, y])

    print(f.forward([-1.36003, -1.35962]))

    print("min:", np.min(z))

    # 显示图像
    ax = fig.add_subplot(111, projection='3d')
    surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.xlabel('x')
    plt.ylabel('y')
    plt.show()
